from transformers import AutoModelForCausalLM
from peft import LoraConfig, get_peft_model, TaskType

# 尝试在 q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj 层添加 Lora 结构，下面使用 PEFT 库实现，这里 r 使用 8，lora_alpha 使用 32
model_path = "/root/ModelSpace/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto")
peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM,
                         target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
                         inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()
print(model)
